package com.eric.study.reactor.filter;

import com.eric.study.reactor.commons.ConstantsFields;
import lombok.extern.slf4j.Slf4j;
import org.slf4j.MDC;
import org.springframework.stereotype.Component;
import org.springframework.web.server.ServerWebExchange;
import org.springframework.web.server.WebFilter;
import org.springframework.web.server.WebFilterChain;
import reactor.core.publisher.Mono;
import reactor.util.context.Context;

import java.util.Optional;
import java.util.UUID;

/**
 * <p> TestFilter.
 *
 * <p>
 * change history:
 *
 * <pre>
 * date         defect       person          comments
 * --------------------------------------------------------------------------------
 * 2023/9/15    ********     Qiheng Zhang    create file.
 * </pre>
 *
 * @author Qiheng Zhang
 * @date 2023/9/15 10:42
 * @copyright 2023, FA Software (Shanghai) Co., Ltd. All Rights Reserved.
 * @since 1.0.0
 */
@Component
@Slf4j
public class TraceIdFilter implements WebFilter {


    @Override
    public Mono<Void> filter(ServerWebExchange exchange, WebFilterChain chain) {
        return chain.filter(exchange)
                // 放入到当前上下文中
                .contextWrite(context -> {
                    String traceId = Optional.ofNullable(exchange.getRequest().getHeaders().getFirst(ConstantsFields.TRACE_ID))
                            .orElse(UUID.randomUUID().toString());
                    MDC.put("traceId", traceId);
                    log.info("生成 trace id = [{}]", traceId);

                    final Context ret = context.put(ConstantsFields.TRACE_ID, traceId);
                    exchange.getAttributes().put(ConstantsFields.TRACE_ID, traceId);
                    return ret;
                });
    }
}
